import os
import pandas as pd
import tempfile
import shutil
from PIL import Image
import io
import chardet
import zipfile
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')


class FileProcessor:
    def __init__(self):
        self.temp_files = []
        
    def __del__(self):
        """清理临时文件"""
        self.cleanup()
    
    def cleanup(self):
        """删除所有临时文件"""
        for temp_file in self.temp_files:
            if os.path.exists(temp_file):
                try:
                    os.remove(temp_file)
                except:
                    pass
        self.temp_files = []
    
    def check_file_corruption(self, file_path):
        """
        检查文件是否损坏[1,2](@ref)
        
        Returns:
            tuple: (是否损坏, 错误信息)
        """
        if not os.path.exists(file_path):
            return True, "文件不存在"
        
        file_size = os.path.getsize(file_path)
        if file_size == 0:
            return True, "文件大小为0"
        
        file_ext = os.path.splitext(file_path)[1].lower()
        
        try:
            if file_ext == '.csv':
                return self._check_csv_corruption(file_path)
            elif file_ext == '.xlsx':
                return self._check_xlsx_corruption(file_path)
            elif file_ext == '.xls':
                return self._check_xls_corruption(file_path)
            else:
                return True, f"不支持的文件格式: {file_ext}"
        except Exception as e:
            return True, f"文件检查错误: {str(e)}"
    
    def _check_csv_corruption(self, file_path):
        """检查CSV文件是否损坏[3](@ref)"""
        try:
            # 检测文件编码
            with open(file_path, 'rb') as f:
                encoding = chardet.detect(f.read())['encoding']
            
            if encoding is None:
                return True, "无法检测文件编码"
            
            # 尝试用pandas读取
            try:
                df = pd.read_csv(file_path, encoding=encoding, engine='python')
                return False, "文件正常"
            except UnicodeDecodeError:
                # 尝试常见编码
                for enc in ['utf-8', 'gbk', 'latin1']:
                    try:
                        df = pd.read_csv(file_path, encoding=enc, engine='python')
                        return False, "文件正常"
                    except:
                        continue
                return True, "CSV文件编码错误或格式损坏"
            except Exception as e:
                return True, f"CSV文件读取错误: {str(e)}"
                
        except Exception as e:
            return True, f"CSV文件检查异常: {str(e)}"
    
    def _check_xlsx_corruption(self, file_path):
        """检查XLSX文件是否损坏[1](@ref)"""
        try:
            # 检查文件头[1](@ref)
            with open(file_path, 'rb') as f:
                header = f.read(4).hex()
                if header != '504b0304':  # ZIP文件头
                    return True, "XLSX文件头错误"
            
            # 检查ZIP文件完整性
            try:
                with zipfile.ZipFile(file_path, 'r') as zip_ref:
                    if not zip_ref.testzip():
                        return False, "文件正常"
                    else:
                        return True, "ZIP压缩包损坏"
            except zipfile.BadZipFile:
                return True, "不是有效的ZIP文件"
                
        except Exception as e:
            return True, f"XLSX文件检查异常: {str(e)}"
    
    def _check_xls_corruption(self, file_path):
        """检查XLS文件是否损坏[2](@ref)"""
        try:
            # 检查文件头
            with open(file_path, 'rb') as f:
                header = f.read(8).hex()
                expected_header = 'd0cf11e0a1b11ae1'  # XLS文件头
                if header != expected_header:
                    return True, "XLS文件头错误"
            
            # 尝试用pandas读取
            try:
                df = pd.read_excel(file_path, engine='xlrd')
                return False, "文件正常"
            except Exception as e:
                return True, f"XLS文件读取错误: {str(e)}"
                
        except Exception as e:
            return True, f"XLS文件检查异常: {str(e)}"
    
    def repair_file(self, file_path):
        """
        尝试修复损坏的文件[1,4](@ref)
        
        Returns:
            tuple: (修复后的文件路径, 是否成功)
        """
        try:
            file_ext = os.path.splitext(file_path)[1].lower()
            temp_file = tempfile.mktemp(suffix=file_ext)
            self.temp_files.append(temp_file)
            
            if file_ext == '.csv':
                return self._repair_csv(file_path, temp_file)
            elif file_ext == '.xlsx':
                return self._repair_xlsx(file_path, temp_file)
            elif file_ext == '.xls':
                return self._repair_xls(file_path, temp_file)
            else:
                return file_path, False
                
        except Exception as e:
            return file_path, False
    
    def _repair_csv(self, file_path, temp_file):
        """修复CSV文件[3](@ref)"""
        try:
            # 检测编码并重新编码为UTF-8
            with open(file_path, 'rb') as f:
                raw_data = f.read()
                encoding = chardet.detect(raw_data)['encoding']
            
            if encoding is None:
                encoding = 'utf-8'
            
            # 尝试不同编码读取
            for enc in [encoding, 'utf-8', 'gbk', 'latin1']:
                try:
                    df = pd.read_csv(file_path, encoding=enc, engine='python', on_bad_lines='skip')
                    df.to_csv(temp_file, index=False, encoding='utf-8')
                    return temp_file, True
                except:
                    continue
            
            # 终极方案：文本提取
            with open(file_path, 'rb') as f:
                content = f.read()
                text_content = content.decode('latin1', errors='ignore')
            
            # 简单的CSV格式处理
            lines = text_content.split('\n')
            cleaned_lines = []
            for line in lines:
                if ',' in line:  # 基本确认是CSV格式
                    cleaned_lines.append(line)
            
            with open(temp_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(cleaned_lines))
            
            return temp_file, True
            
        except Exception as e:
            return file_path, False
    
    def _repair_xlsx(self, file_path, temp_file):
        """修复XLSX文件[1](@ref)"""
        try:
            # 方法1: 使用pandas跳过错误行
            try:
                df = pd.read_excel(file_path, engine='openpyxl', on_bad_lines='skip')
                df.to_excel(temp_file, index=False)
                return temp_file, True
            except:
                pass
            
            # 方法2: 尝试用openpyxl
            try:
                from openpyxl import load_workbook
                wb = load_workbook(file_path, read_only=True, data_only=True)
                df = pd.DataFrame()
                for sheet_name in wb.sheetnames:
                    sheet = wb[sheet_name]
                    data = []
                    for row in sheet.iter_rows(values_only=True):
                        data.append(row)
                    if data:
                        sheet_df = pd.DataFrame(data)
                        df = pd.concat([df, sheet_df], ignore_index=True)
                df.to_excel(temp_file, index=False)
                return temp_file, True
            except:
                pass
            
            return file_path, False
            
        except Exception as e:
            return file_path, False
    
    def _repair_xls(self, file_path, temp_file):
        """修复XLS文件[4](@ref)"""
        try:
            # 方法1: 使用pandas的不同引擎
            for engine in ['xlrd', 'calamine']:
                try:
                    df = pd.read_excel(file_path, engine=engine)
                    df.to_excel(temp_file, index=False)
                    return temp_file, True
                except:
                    continue
            
            # 方法2: 尝试使用win32com（Windows系统）[4](@ref)
            try:
                import win32com.client
                excel = win32com.client.Dispatch("Excel.Application")
                excel.Visible = False
                wb = excel.Workbooks.Open(os.path.abspath(file_path))
                wb.SaveAs(os.path.abspath(temp_file), FileFormat=51)  # xlsx格式
                wb.Close()
                excel.Quit()
                return temp_file, True
            except:
                pass
            
            return file_path, False
            
        except Exception as e:
            return file_path, False
    
    def parse_file_content(self, file_path):
        """
        解析文件内容，提取文本、图片、表格和超链接[6,7](@ref)
        
        Returns:
            dict: 包含解析结果的字典
        """
        result = {
            'text': '',
            'images': [],
            'tables': '',
            'hyperlinks': []
        }
        
        file_ext = os.path.splitext(file_path)[1].lower()
        
        try:
            if file_ext == '.csv':
                return self._parse_csv(file_path)
            elif file_ext in ['.xlsx', '.xls']:
                return self._parse_excel(file_path)
            else:
                return result
        except Exception as e:
            print(f"解析错误: {e}")
            return result
    
    def _parse_csv(self, file_path):
        """解析CSV文件[6](@ref)"""
        result = {
            'text': '',
            'images': [],
            'tables': '',
            'hyperlinks': []
        }
        
        try:
            # 检测编码
            with open(file_path, 'rb') as f:
                encoding = chardet.detect(f.read())['encoding']
            
            if encoding is None:
                encoding = 'utf-8'
            
            # 读取CSV
            for enc in [encoding, 'utf-8', 'gbk', 'latin1']:
                try:
                    df = pd.read_csv(file_path, encoding=enc, engine='python')
                    break
                except:
                    continue
            else:
                # 所有编码都失败，尝试原始读取
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
                result['text'] = content
                return result
            
            # 转换为markdown表格[6,7](@ref)
            if not df.empty:
                result['tables'] = df.to_markdown(index=False)
                result['text'] = f"CSV文件包含 {len(df)} 行数据\n\n"
            
        except Exception as e:
            print(f"CSV解析错误: {e}")
        
        return result
    
    def _parse_excel(self, file_path):
        """解析Excel文件[6,7](@ref)"""
        result = {
            'text': '',
            'images': [],
            'tables': '',
            'hyperlinks': []
        }
        
        try:
            # 读取Excel文件
            file_ext = os.path.splitext(file_path)[1].lower()
            
            if file_ext == '.xlsx':
                engine = 'openpyxl'
            else:
                engine = 'xlrd'
            
            # 获取所有sheet名
            excel_file = pd.ExcelFile(file_path, engine=engine)
            sheet_names = excel_file.sheet_names
            
            all_tables = []
            
            for sheet_name in sheet_names:
                try:
                    df = pd.read_excel(file_path, sheet_name=sheet_name, engine=engine, header=None)
                    
                    if not df.empty:
                        # 清理NaN值
                        df = df.fillna('')
                        
                        # 转换为markdown[7](@ref)
                        table_md = f"### {sheet_name}\n\n"
                        table_md += df.to_markdown(index=False, header=False)
                        table_md += "\n\n"
                        all_tables.append(table_md)
                        
                except Exception as e:
                    print(f"解析sheet {sheet_name} 错误: {e}")
                    continue
            
            result['tables'] = '\n'.join(all_tables)
            result['text'] = f"Excel文件包含 {len(sheet_names)} 个工作表\n\n"
            
            # 尝试提取图片（仅xlsx支持）[7](@ref)
            if file_ext == '.xlsx':
                try:
                    from openpyxl import load_workbook
                    wb = load_workbook(file_path)
                    
                    for sheet_name in wb.sheetnames:
                        sheet = wb[sheet_name]
                        for image in sheet._images:
                            if hasattr(image, 'ref'):
                                try:
                                    img_data = image._data()
                                    img = Image.open(io.BytesIO(img_data))
                                    result['images'].append(img)
                                except Exception as e:
                                    print(f"图片提取错误: {e}")
                                    
                except Exception as e:
                    print(f"图片提取失败: {e}")
            
            # 提取超链接（仅xlsx支持）
            if file_ext == '.xlsx':
                try:
                    from openpyxl import load_workbook
                    wb = load_workbook(file_path)
                    
                    for sheet_name in wb.sheetnames:
                        sheet = wb[sheet_name]
                        for row in sheet.iter_rows():
                            for cell in row:
                                if cell.hyperlink and cell.hyperlink.target:
                                    result['hyperlinks'].append(cell.hyperlink.target)
                                    
                except Exception as e:
                    print(f"超链接提取失败: {e}")
                    
        except Exception as e:
            print(f"Excel解析错误: {e}")
        
        return result
    
    def process_file(self, file_path):
        """
        处理文件的完整流程[1,2,3](@ref)
        
        Returns:
            dict: 处理结果
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            return {"error": "文件不存在"}
        
        # 检查文件是否损坏
        is_corrupted, error_msg = self.check_file_corruption(file_path)
        
        working_file = file_path
        repaired = False
        
        if is_corrupted:
            print(f"文件可能损坏: {error_msg}, 尝试修复...")
            working_file, repaired = self.repair_file(file_path)
            
            if not repaired:
                return {"error": f"文件损坏且修复失败: {error_msg}"}
        
        # 解析文件内容
        result = self.parse_file_content(working_file)
        result['repaired'] = repaired
        result['original_file'] = file_path
        
        # 组合最终的markdown文本
        markdown_content = []
        
        if result['text']:
            markdown_content.append(result['text'])
        
        if result['tables']:
            markdown_content.append("## 表格数据\n")
            markdown_content.append(result['tables'])
        
        if result['hyperlinks']:
            markdown_content.append("## 超链接\n")
            for link in result['hyperlinks']:
                markdown_content.append(f"- {link}")
        
        result['markdown'] = '\n'.join(markdown_content)
        
        return result


# 使用示例
def parse_table(file_path):
    processor = FileProcessor()
    
    try:
        # 处理文件
        result = processor.process_file(file_path)
        
        if 'error' in result:
            print(f"处理失败: {result['error']}")
            return None
        else:
            print("=== 处理结果 ===")
            print(f"文件是否修复: {result['repaired']}")
            print(f"提取图片数量: {len(result['images'])}")
            print(f"提取超链接数量: {len(result['hyperlinks'])}")
            print("\n=== Markdown内容 ===")
            print(result['markdown'][:1000] + "..." if len(result['markdown']) > 1000 else result['markdown'])
            return result
    finally:
        # 清理临时文件
        processor.cleanup()
        return None


if __name__ == "__main__":
    parse_table("path/example.xlsx")